from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np
import random
from .padding_resize_convert import PaddingResizeConvert


class SizeChangeConvert(ConvertBase):
    
    def __init__(self, target_size, clip_size = None, offset = (10, 10), min_scale = 0.5, max_scale = 2, is_random =True, use_rate = 0.5, filter_small_box_rate = 0.5):
        super().__init__(use_rate)
        
        self.target_size = target_size
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.is_rand = is_random
        self.scale = 1
        
        self.filter_small_box_rate = filter_small_box_rate # 过滤box较小的阈值
        
        self.offset = offset
        
        if clip_size is None:
            self.clip_size = self.target_size
        else:
            self.clip_size = clip_size
        
        # self.padding_resize = PaddingResizeConvert(self.target_size)
        
    def convert(self, img, boxes, points):
        
        if self.is_rand:
            self.scale = self.min_scale + (self.max_scale - self.min_scale) * random.random()
        
        img, boxes, points = PaddingResizeConvert((int(self.target_size[0] * self.scale), int(self.target_size[1] * self.scale))).convert(img, boxes, points)
        
        offset_xy = np.array([random.randint(0, self.offset[0]), random.randint(0, self.offset[1])], np.float32)
        r = random.randint(0, 1)
        if r == 0:
            min_xy = (np.min(boxes[:, [0, 2]]), np.min(boxes[:, [1, 3]]))
            max_xy = (min_xy[0] + self.clip_size[0], min_xy[1] + self.clip_size[1])
        else:
            max_xy = (np.max(boxes[:, [0, 2]]), np.max(boxes[:, [1, 3]]))
            min_xy = (max_xy[0] - self.clip_size[0], max_xy[1] - self.clip_size[1])    
        rang = np.array([min_xy[0], min_xy[1], max_xy[0], max_xy[1]])
        rang = (rang.reshape([-1, 2]) + offset_xy).reshape([4])
        
        rang = np.clip(rang, 0, np.max(img.shape)).astype(np.int32)
        img = img[rang[1]:rang[3], rang[0] : rang[2]]
        
        if boxes is not None:
            res_box = []
            for box in boxes:
                new_box = [box[0] if box[0] > rang[0] else rang[0], 
                        box[1] if box[1] > rang[1] else rang[1],
                        box[2] if box[2] < rang[2] else rang[2],
                        box[3] if box[3] < rang[3] else rang[3]]
                new_box = (np.array(new_box).reshape([-1, 2]) - rang[[0, 1]]).reshape([4])
                box_area = (box[2] - box[0]) * (box[3] - box[1])
                new_box_area = (new_box[2] - new_box[0]) * (new_box[3] - new_box[1])

                if new_box_area / box_area < self.filter_small_box_rate:
                    continue
                    
                else:
                    res_box.append(new_box)
            
            boxes = np.array(res_box)        
            
        if len(boxes) == 0:
            return None, None, None
        img, boxes, points = PaddingResizeConvert(self.target_size).convert(img, boxes, points)
        
        return img, boxes, points
        
        
        
        